// This is an open source non-commercial project. Dear PVS-Studio, please check it.
// PVS-Studio Static Code Analyzer for C, C++ and C#: http://www.viva64.com

// ReSharper disable CheckNamespace
// ReSharper disable CommentTypo
// ReSharper disable IdentifierTypo
// ReSharper disable InconsistentNaming
// ReSharper disable UnusedMember.Global

/* PdfStandardSecurityHandler.cs --
 * Ars Magna project, http://arsmagna.ru
 */

#region Using directives

using System;
using System.Diagnostics;
using System.Security.Cryptography;

using AM;

using PdfSharpCore.Pdf.IO;
using PdfSharpCore.Pdf.Internal;

#endregion

#nullable enable

#pragma warning disable 0169
#pragma warning disable 0649

namespace PdfSharpCore.Pdf.Security;

/// <summary>
/// Represents the standard PDF security handler.
/// </summary>
public sealed class PdfStandardSecurityHandler : PdfSecurityHandler
{
    #region Construction

    internal PdfStandardSecurityHandler (PdfDocument document)
        : base (document)
    {
        // пустое тело конструктора
    }

    internal PdfStandardSecurityHandler (PdfDictionary dict)
        : base (dict)
    {
        // пустое тело конструктора
    }

    #endregion

    /// <summary>
    /// Sets the user password of the document. Setting a password automatically sets the
    /// PdfDocumentSecurityLevel to PdfDocumentSecurityLevel.Encrypted128Bit if its current
    /// value is PdfDocumentSecurityLevel.None.
    /// </summary>
    public string UserPassword
    {
        set
        {
            if (_document!._securitySettings.DocumentSecurityLevel == PdfDocumentSecurityLevel.None)
            {
                _document._securitySettings.DocumentSecurityLevel = PdfDocumentSecurityLevel.Encrypted128Bit;
            }

            _userPassword = value;
        }
    }

    internal string? _userPassword;

    /// <summary>
    /// Sets the owner password of the document. Setting a password automatically sets the
    /// PdfDocumentSecurityLevel to PdfDocumentSecurityLevel.Encrypted128Bit if its current
    /// value is PdfDocumentSecurityLevel.None.
    /// </summary>
    public string OwnerPassword
    {
        set
        {
            if (_document!._securitySettings.DocumentSecurityLevel == PdfDocumentSecurityLevel.None)
            {
                _document._securitySettings.DocumentSecurityLevel = PdfDocumentSecurityLevel.Encrypted128Bit;
            }

            _ownerPassword = value;
        }
    }

    internal string? _ownerPassword;

    /// <summary>
    /// Gets or sets the user access permission represented as an integer in the P key.
    /// </summary>
    internal PdfUserAccessPermission Permission
    {
        get
        {
            var permission = (PdfUserAccessPermission)Elements.GetInteger (Keys.P);
            if ((int)permission == 0)
            {
                permission = PdfUserAccessPermission.PermitAll;
            }

            return permission;
        }
        set => Elements.SetInteger (Keys.P, (int)value);
    }

    /// <summary>
    /// Encrypts the whole document.
    /// </summary>
    public void EncryptDocument()
    {
        if (_document is not null)
        {
            foreach (var iref in _document._irefTable.AllReferences)
            {
                if (!ReferenceEquals (iref.Value, this))
                {
                    EncryptObject (iref.Value);
                }
            }
        }
    }

    /// <summary>
    /// Encrypts an indirect object.
    /// </summary>
    internal void EncryptObject (PdfObject value)
    {
        Debug.Assert (value.Reference != null);

        SetHashKey (value.ObjectID);
#if DEBUG
        if (value.ObjectID.ObjectNumber == 10)
        {
            GetType().NotUsed();
        }
#endif

        if (value is PdfDictionary dict)
        {
            EncryptDictionary (dict);
        }
        else if (value is PdfArray array)
        {
            EncryptArray (array);
        }
        else if (value is PdfStringObject { Length: not 0 } str)
        {
            var bytes = str.EncryptionValue;
            PrepareKey();
            EncryptRC4 (bytes);
            str.EncryptionValue = bytes;
        }
    }

    /// <summary>
    /// Encrypts a dictionary.
    /// </summary>
    void EncryptDictionary (PdfDictionary dict)
    {
        var names = dict.Elements.KeyNames;
        foreach (var item in dict.Elements)
        {
            switch (item.Value)
            {
                case PdfString value1:
                    EncryptString (value1);
                    break;

                case PdfDictionary value2:
                    EncryptDictionary (value2);
                    break;

                case PdfArray value3:
                    EncryptArray (value3);
                    break;
            }
        }

        if (dict.Stream != null)
        {
            var bytes = dict.Stream.Value;
            if (bytes.Length != 0)
            {
                PrepareKey();
                EncryptRC4 (bytes);
                dict.Stream.Value = bytes;
            }
        }
    }

    /// <summary>
    /// Encrypts an array.
    /// </summary>
    void EncryptArray (PdfArray array)
    {
        var count = array.Elements.Count;
        for (var idx = 0; idx < count; idx++)
        {
            var item = array.Elements[idx];
            switch (item)
            {
                case PdfString value1:
                    EncryptString (value1);
                    break;

                case PdfDictionary value2:
                    EncryptDictionary (value2);
                    break;

                case PdfArray value3:
                    EncryptArray (value3);
                    break;
            }
        }
    }

    /// <summary>
    /// Encrypts a string.
    /// </summary>
    void EncryptString (PdfString value)
    {
        if (value.Length != 0)
        {
            var bytes = value.EncryptionValue;
            PrepareKey();
            EncryptRC4 (bytes);
            value.EncryptionValue = bytes;
        }
    }

    /// <summary>
    /// Encrypts an array.
    /// </summary>
    internal byte[]? EncryptBytes (byte[]? bytes)
    {
        if (bytes != null && bytes.Length != 0)
        {
            PrepareKey();
            EncryptRC4 (bytes);
        }

        return bytes;
    }

    #region Encryption Algorithms

    /// <summary>
    /// Checks the password.
    /// </summary>
    /// <param name="inputPassword">Password or null if no password is provided.</param>
    public PasswordValidity ValidatePassword (string? inputPassword)
    {
        // We can handle 40 and 128 bit standard encryption.
        var filter = Elements.GetName (PdfSecurityHandler.Keys.Filter);
        var v = Elements.GetInteger (PdfSecurityHandler.Keys.V);
        if (filter != "/Standard" || !(v >= 1 && v <= 3))
        {
            throw new PdfReaderException (PSSR.UnknownEncryption);
        }

        var owner = Owner.ThrowIfNull();
        var document = _document.ThrowIfNull();

        var documentId = PdfEncoders.RawEncoding.GetBytes (owner.Internals.FirstDocumentID);
        var oValue = PdfEncoders.RawEncoding.GetBytes (Elements.GetString (Keys.O));
        var uValue = PdfEncoders.RawEncoding.GetBytes (Elements.GetString (Keys.U));
        var pValue = Elements.GetInteger (Keys.P);
        var rValue = Elements.GetInteger (Keys.R);

        inputPassword ??= string.Empty;

        var strongEncryption = rValue == 3;
        var keyLength = strongEncryption ? 16 : 32;

        // Try owner password first.
        //byte[] password = PdfEncoders.RawEncoding.GetBytes(inputPassword);
        InitWithOwnerPassword (documentId, inputPassword, oValue, pValue, strongEncryption);
        if (EqualsKey (uValue, keyLength))
        {
            document.SecuritySettings.HasOwnerPermissions = true;
            return PasswordValidity.OwnerPassword;
        }

        document.SecuritySettings.HasOwnerPermissions = false;

        // Now try user password.
        //password = PdfEncoders.RawEncoding.GetBytes(inputPassword);
        InitWithUserPassword (documentId, inputPassword, oValue, pValue, strongEncryption);
        if (EqualsKey (uValue, keyLength))
        {
            return PasswordValidity.UserPassword;
        }

        return PasswordValidity.Invalid;
    }

    [Conditional ("DEBUG")]
    static void DumpBytes (string tag, byte[] bytes)
    {
        var dump = tag + ": ";
        for (var idx = 0; idx < bytes.Length; idx++)
        {
            dump += $"{bytes[idx]:X2}";
        }

        Debug.WriteLine (dump);
    }

    /// <summary>
    /// Pads a password to a 32 byte array.
    /// </summary>
    static byte[] PadPassword (string? password)
    {
        var padded = new byte[32];
        if (password == null)
        {
            Array.Copy (PasswordPadding, 0, padded, 0, 32);
        }
        else
        {
            var length = password.Length;
            Array.Copy (PdfEncoders.RawEncoding.GetBytes (password), 0, padded, 0, Math.Min (length, 32));
            if (length < 32)
            {
                Array.Copy (PasswordPadding, 0, padded, length, 32 - length);
            }
        }

        return padded;
    }

    static readonly byte[] PasswordPadding = // 32 bytes password padding defined by Adobe
    {
        0x28, 0xBF, 0x4E, 0x5E, 0x4E, 0x75, 0x8A, 0x41, 0x64, 0x00, 0x4E, 0x56, 0xFF, 0xFA, 0x01, 0x08,
        0x2E, 0x2E, 0x00, 0xB6, 0xD0, 0x68, 0x3E, 0x80, 0x2F, 0x0C, 0xA9, 0xFE, 0x64, 0x53, 0x69, 0x7A,
    };

    /// <summary>
    /// Generates the user key based on the padded user password.
    /// </summary>
    void InitWithUserPassword (byte[] documentId, string userPassword, byte[] ownerKey, int permissions,
        bool strongEncryption)
    {
        InitEncryptionKey (documentId, PadPassword (userPassword), ownerKey, permissions, strongEncryption);
        SetupUserKey (documentId);
    }

    /// <summary>
    /// Generates the user key based on the padded owner password.
    /// </summary>
    void InitWithOwnerPassword (byte[] documentId, string ownerPassword, byte[] ownerKey, int permissions,
        bool strongEncryption)
    {
        var userPad = ComputeOwnerKey (ownerKey, PadPassword (ownerPassword), strongEncryption);
        InitEncryptionKey (documentId, userPad, ownerKey, permissions, strongEncryption);
        SetupUserKey (documentId);
    }

    /// <summary>
    /// Computes the padded user password from the padded owner password.
    /// </summary>
    byte[] ComputeOwnerKey (byte[] userPad, byte[] ownerPad, bool strongEncryption)
    {
        var ownerKey = new byte[32];

        //#if !SILVERLIGHT
        var digest = _md5.ComputeHash (ownerPad);
        if (strongEncryption)
        {
            var mkey = new byte[16];

            // Hash the pad 50 times
            for (var idx = 0; idx < 50; idx++)
            {
                digest = _md5.ComputeHash (digest);
            }

            Array.Copy (userPad, 0, ownerKey, 0, 32);

            // Encrypt the key
            for (var i = 0; i < 20; i++)
            {
                for (var j = 0; j < mkey.Length; ++j)
                {
                    mkey[j] = (byte)(digest[j] ^ i);
                }

                PrepareRC4Key (mkey);
                EncryptRC4 (ownerKey);
            }
        }
        else
        {
            PrepareRC4Key (digest, 0, 5);
            EncryptRC4 (userPad, ownerKey);
        }

        //#endif
        return ownerKey;
    }

    /// <summary>
    /// Computes the encryption key.
    /// </summary>
    void InitEncryptionKey (byte[] documentID, byte[] userPad, byte[] ownerKey, int permissions, bool strongEncryption)
    {
        _ownerKey = ownerKey;
        _encryptionKey = new byte[strongEncryption ? 16 : 5];

        _md5.Initialize();
        _md5.TransformBlock (userPad, 0, userPad.Length, userPad, 0);
        _md5.TransformBlock (ownerKey, 0, ownerKey.Length, ownerKey, 0);

        // Split permission into 4 bytes
        var permission = new byte[4];
        permission[0] = (byte)permissions;
        permission[1] = (byte)(permissions >> 8);
        permission[2] = (byte)(permissions >> 16);
        permission[3] = (byte)(permissions >> 24);

        _md5.TransformBlock (permission, 0, 4, permission, 0);
        _md5.TransformBlock (documentID, 0, documentID.Length, documentID, 0);
        _md5.TransformFinalBlock (permission, 0, 0);
        var digest = _md5.Hash;
        _md5.Initialize();

        // Create the hash 50 times (only for 128 bit)
        if (_encryptionKey.Length == 16)
        {
            for (var idx = 0; idx < 50; idx++)
            {
                digest = _md5.ComputeHash (digest!);
                _md5.Initialize();
            }
        }

        Array.Copy (digest!, 0, _encryptionKey, 0, _encryptionKey.Length);
    }

    /// <summary>
    /// Computes the user key.
    /// </summary>
    void SetupUserKey (byte[] documentId)
    {
        if (_encryptionKey!.Length == 16)
        {
            _md5.TransformBlock (PasswordPadding, 0, PasswordPadding.Length, PasswordPadding, 0);
            _md5.TransformFinalBlock (documentId, 0, documentId.Length);
            var digest = _md5.Hash;
            _md5.Initialize();
            Array.Copy (digest!, 0, _userKey, 0, 16);
            for (var idx = 16; idx < 32; idx++)
            {
                _userKey[idx] = 0;
            }

            //Encrypt the key
            for (var i = 0; i < 20; i++)
            {
                for (var j = 0; j < _encryptionKey.Length; j++)
                {
                    digest![j] = (byte)(_encryptionKey[j] ^ i);
                }

                PrepareRC4Key (digest, 0, _encryptionKey.Length);
                EncryptRC4 (_userKey, 0, 16);
            }
        }
        else
        {
            PrepareRC4Key (_encryptionKey);
            EncryptRC4 (PasswordPadding, _userKey);
        }
    }

    /// <summary>
    /// Prepare the encryption key.
    /// </summary>
    void PrepareKey()
    {
        PrepareRC4Key (_key!, 0, _keySize);
    }

    /// <summary>
    /// Prepare the encryption key.
    /// </summary>
    void PrepareRC4Key (byte[] key)
    {
        PrepareRC4Key (key, 0, key.Length);
    }

    /// <summary>
    /// Prepare the encryption key.
    /// </summary>
    void PrepareRC4Key (byte[] key, int offset, int length)
    {
        var idx1 = 0;
        var idx2 = 0;
        for (var idx = 0; idx < 256; idx++)
        {
            _state[idx] = (byte)idx;
        }

        byte tmp;
        for (var idx = 0; idx < 256; idx++)
        {
            idx2 = (key[idx1 + offset] + _state[idx] + idx2) & 255;
            tmp = _state[idx];
            _state[idx] = _state[idx2];
            _state[idx2] = tmp;
            idx1 = (idx1 + 1) % length;
        }
    }

    /// <summary>
    /// Encrypts the data.
    /// </summary>

    // ReSharper disable InconsistentNaming
    void EncryptRC4 (byte[] data)

        // ReSharper restore InconsistentNaming
    {
        EncryptRC4 (data, 0, data.Length, data);
    }

    /// <summary>
    /// Encrypts the data.
    /// </summary>

    // ReSharper disable once InconsistentNaming
    void EncryptRC4 (byte[] data, int offset, int length)
    {
        EncryptRC4 (data, offset, length, data);
    }

    /// <summary>
    /// Encrypts the data.
    /// </summary>

    // ReSharper disable once InconsistentNaming
    void EncryptRC4 (byte[] inputData, byte[] outputData)
    {
        EncryptRC4 (inputData, 0, inputData.Length, outputData);
    }

    /// <summary>
    /// Encrypts the data.
    /// </summary>

    // ReSharper disable once InconsistentNaming
    void EncryptRC4 (byte[] inputData, int offset, int length, byte[] outputData)
    {
        length += offset;
        int x = 0, y = 0;
        byte b;
        for (var idx = offset; idx < length; idx++)
        {
            x = (x + 1) & 255;
            y = (_state[x] + y) & 255;
            b = _state[x];
            _state[x] = _state[y];
            _state[y] = b;
            outputData[idx] = (byte)(inputData[idx] ^ _state[(_state[x] + _state[y]) & 255]);
        }
    }

    /// <summary>
    /// Checks whether the calculated key correct.
    /// </summary>
    bool EqualsKey (byte[] value, int length)
    {
        for (var idx = 0; idx < length; idx++)
        {
            if (_userKey[idx] != value[idx])
            {
                return false;
            }
        }

        return true;
    }

    /// <summary>
    /// Set the hash key for the specified object.
    /// </summary>
    internal void SetHashKey (PdfObjectID id)
    {
        var objectId = new byte[5];
        _md5.Initialize();

        // Split the object number and generation
        objectId[0] = (byte)id.ObjectNumber;
        objectId[1] = (byte)(id.ObjectNumber >> 8);
        objectId[2] = (byte)(id.ObjectNumber >> 16);
        objectId[3] = (byte)id.GenerationNumber;
        objectId[4] = (byte)(id.GenerationNumber >> 8);
        _md5.TransformBlock (_encryptionKey!, 0, _encryptionKey!.Length, _encryptionKey, 0);
        _md5.TransformFinalBlock (objectId, 0, objectId.Length);
        _key = _md5.Hash;
        _md5.Initialize();
        _keySize = _encryptionKey.Length + 5;
        if (_keySize > 16)
        {
            _keySize = 16;
        }
    }

    /// <summary>
    /// Prepares the security handler for encrypting the document.
    /// </summary>
    public void PrepareEncryption()
    {
        //#if !SILVERLIGHT
        Debug.Assert (_document!._securitySettings.DocumentSecurityLevel != PdfDocumentSecurityLevel.None);
        var permissions = (int)Permission;
        var strongEncryption = _document._securitySettings.DocumentSecurityLevel ==
                               PdfDocumentSecurityLevel.Encrypted128Bit;

        PdfInteger vValue;
        PdfInteger length;
        PdfInteger rValue;

        if (strongEncryption)
        {
            vValue = new PdfInteger (2);
            length = new PdfInteger (128);
            rValue = new PdfInteger (3);
        }
        else
        {
            vValue = new PdfInteger (1);
            length = new PdfInteger (40);
            rValue = new PdfInteger (2);
        }

        if (string.IsNullOrEmpty (_userPassword))
        {
            _userPassword = "";
        }

        // Use user password twice if no owner password provided.
        if (string.IsNullOrEmpty (_ownerPassword))
        {
            _ownerPassword = _userPassword;
        }

        // Correct permission bits
        permissions |= (int)(strongEncryption ? (uint)0xfffff0c0 : (uint)0xffffffc0);
        permissions &= unchecked ((int)0xfffffffc);

        var pValue = new PdfInteger (permissions);

        Debug.Assert (_ownerPassword.Length > 0, "Empty owner password.");
        var userPad = PadPassword (_userPassword);
        var ownerPad = PadPassword (_ownerPassword);

        _md5.Initialize();
        _ownerKey = ComputeOwnerKey (userPad, ownerPad, strongEncryption);
        var documentID = PdfEncoders.RawEncoding.GetBytes (_document.Internals.FirstDocumentID);
        InitWithUserPassword (documentID, _userPassword, _ownerKey, permissions, strongEncryption);

        var oValue = new PdfString (PdfEncoders.RawEncoding.GetString (_ownerKey, 0, _ownerKey.Length));
        var uValue = new PdfString (PdfEncoders.RawEncoding.GetString (_userKey, 0, _userKey.Length));

        Elements[Keys.Filter] = new PdfName ("/Standard");
        Elements[Keys.V] = vValue;
        Elements[Keys.Length] = length;
        Elements[Keys.R] = rValue;
        Elements[Keys.O] = oValue;
        Elements[Keys.U] = uValue;
        Elements[Keys.P] = pValue;

        //#endif
    }

    /// <summary>
    /// The global encryption key.
    /// </summary>
    private byte[]? _encryptionKey;

    readonly MD5 _md5 = MD5.Create();

    /// <summary>
    /// Bytes used for RC4 encryption.
    /// </summary>
    readonly byte[] _state = new byte[256];

    /// <summary>
    /// The encryption key for the owner.
    /// </summary>
    byte[] _ownerKey = new byte[32];

    /// <summary>
    /// The encryption key for the user.
    /// </summary>
    readonly byte[] _userKey = new byte[32];

    /// <summary>
    /// The encryption key for a particular object/generation.
    /// </summary>
    private byte[]? _key;

    /// <summary>
    /// The encryption key length for a particular object/generation.
    /// </summary>
    int _keySize;

    #endregion

    internal override void WriteObject (PdfWriter writer)
    {
        // Don't encrypt myself.
        var securityHandler = writer.SecurityHandler;
        writer.SecurityHandler = null;
        base.WriteObject (writer);
        writer.SecurityHandler = securityHandler;
    }

    #region Keys

    /// <summary>
    /// Predefined keys of this dictionary.
    /// </summary>
    internal new sealed class Keys
        : PdfSecurityHandler.Keys
    {
        /// <summary>
        /// (Required) A number specifying which revision of the standard security handler
        /// should be used to interpret this dictionary:
        /// � 2 if the document is encrypted with a V value less than 2 and does not have any of
        ///   the access permissions set (by means of the P entry, below) that are designated
        ///   "Revision 3 or greater".
        /// � 3 if the document is encrypted with a V value of 2 or 3, or has any "Revision 3 or
        ///   greater" access permissions set.
        /// � 4 if the document is encrypted with a V value of 4
        /// </summary>
        [KeyInfo (KeyType.Integer | KeyType.Required)]
        public const string R = "/R";

        /// <summary>
        /// (Required) A 32-byte string, based on both the owner and user passwords, that is
        /// used in computing the encryption key and in determining whether a valid owner
        /// password was entered.
        /// </summary>
        [KeyInfo (KeyType.String | KeyType.Required)]
        public const string O = "/O";

        /// <summary>
        /// (Required) A 32-byte string, based on the user password, that is used in determining
        /// whether to prompt the user for a password and, if so, whether a valid user or owner
        /// password was entered.
        /// </summary>
        [KeyInfo (KeyType.String | KeyType.Required)]
        public const string U = "/U";

        /// <summary>
        /// (Required) A set of flags specifying which operations are permitted when the document
        /// is opened with user access.
        /// </summary>
        [KeyInfo (KeyType.Integer | KeyType.Required)]
        public const string P = "/P";

        /// <summary>
        /// (Optional; meaningful only when the value of V is 4; PDF 1.5) Indicates whether
        /// the document-level metadata stream is to be encrypted. Applications should respect this value.
        /// Default value: true.
        /// </summary>
        [KeyInfo (KeyType.Boolean | KeyType.Optional)]
        public const string EncryptMetadata = "/EncryptMetadata";

        /// <summary>
        /// Gets the KeysMeta for these keys.
        /// </summary>
        public static DictionaryMeta Meta
        {
            get { return _meta ??= CreateMeta (typeof (Keys)); }
        }

        static DictionaryMeta? _meta;
    }

    /// <summary>
    /// Gets the KeysMeta of this dictionary type.
    /// </summary>
    internal override DictionaryMeta Meta => Keys.Meta;

    #endregion
}
